# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0

# DeepSpeed Team

import deepspeed.comm as dist
from deepspeed.accelerator import get_accelerator

GLOBAL_BATCH_SIZE = 32
MICRO_BATCH_SIZE = 4


def get_train_ds_config(offload, # 是否使用Zero Offload技术
                        dtype, # 训练时前向和反向精度
                        stage=2, # 表示Zero的阶段。Zero有三个阶段（0，1，2）
                        enable_hybrid_engine=False, # 混合引擎。混合引擎可以在训练时进行并行计算：训练和推理一起进行
                        inference_tp_size=1, # 推理期间的流水线划分大小
                        release_inference_cache=False, # 是否每次推理后释放内存
                        pin_parameters=True, # 是否固定参数。固定参数可以避免数据在CPU和GPU之间来回传输。
                        tp_gather_partition_size=8, # 表示Gather操作划分的大小
                        max_out_tokens=512, # 表示最大输出Token
                        enable_tensorboard=False, # 表示是否启用可视化界面
                        enable_mixed_precision_lora=False, # 表示Lora是否启用混合精度
                        tb_path="",
                        tb_name=""):

    device = "cpu" if offload else "none"
    if dtype == "fp16":
        data_type = "fp16"
        dtype_config = {"enabled": True, "loss_scale_window": 100}
    elif dtype == "bf16":
        data_type = "bfloat16"
        dtype_config = {"enabled": True}
    # Zero配置参数
    zero_opt_dict = {
        "stage": stage,
        "overlap_comm": True, # 在梯度通信和反向传播计算之间，是否进行重叠。即有计算完梯度的GPU就开始与其他GPU通信
        "offload_param": {
            "device": device
        },
        "offload_optimizer": {
            "device": device
        },
        "stage3_param_persistence_threshold": 1e4,
        "stage3_max_live_parameters": 3e7,
        "stage3_prefetch_bucket_size": 3e7,
        "memory_efficient_linear": False
    }
    # 混合精度
    if enable_mixed_precision_lora:
        zero_opt_dict["zero_quantized_nontrainable_weights"] = True # 是否对非可训练的权重进行量化
        '''
            # ! dist.get_world_size() 获取当前分布式训练的总进程数
                单节点多卡, word_size=GPU数量
                多节点分布式, word_size=节点数 * GPU数量
            # !  get_accelerator().device_count() 获取当前节点上的可用GPU数量
        '''
        if dist.get_world_size() != get_accelerator().device_count():
            zero_opt_dict["zero_hpz_partition_size"] = get_accelerator(
            ).device_count()
    return {
        "train_batch_size": GLOBAL_BATCH_SIZE, # 全局批量
        "train_micro_batch_size_per_gpu": MICRO_BATCH_SIZE, # 每个GPU上的小批量
        "steps_per_print": 10, # 多少步输出
        "zero_optimization": zero_opt_dict, # Zero优化器
        data_type: dtype_config, # 训练精度
        "gradient_clipping": 1.0, # 梯度裁剪。梯度裁剪阈值，用于防治梯度爆炸
        "prescale_gradients": False, # 是否对梯度进行预缩放
        "wall_clock_breakdown": False, # 是否打印训练过程中各个部分的耗时详细信息
        "hybrid_engine": { # ! 混合引擎的具体作用是什么？ 视频回答：HE允许模型的不同部分在不同的GPU上进行训练
            "enabled": enable_hybrid_engine,
            "max_out_tokens": max_out_tokens,
            "inference_tp_size": inference_tp_size,
            "release_inference_cache": release_inference_cache,
            "pin_parameters": pin_parameters,
            "tp_gather_partition_size": tp_gather_partition_size,
        },
        "tensorboard": { # 训练可视化
            "enabled": enable_tensorboard,
            "output_path": f"{tb_path}/ds_tensorboard_logs/",
            "job_name": f"{tb_name}_tensorboard"
        }
    }


def get_eval_ds_config(offload, dtype, stage=0):
    device = "cpu" if offload else "none"
    if dtype == "fp16":
        data_type = "fp16"
        dtype_config = {
            "enabled": True,
        }
    elif dtype == "bf16":
        data_type = "bfloat16"
        dtype_config = {"enabled": True}
    zero_opt_dict = {
        "stage": stage,
        "stage3_param_persistence_threshold": 1e4,
        "offload_param": {
            "device": device
        },
        "memory_efficient_linear": False
    }
    return {
        "train_batch_size": GLOBAL_BATCH_SIZE,
        "train_micro_batch_size_per_gpu": MICRO_BATCH_SIZE,
        "steps_per_print": 10,
        "zero_optimization": zero_opt_dict,
        data_type: dtype_config,
        "gradient_clipping": 1.0,
        "prescale_gradients": False,
        "wall_clock_breakdown": False
    }
